# 抽象的DFT矩阵分解基类
import pickle
import numpy as np
from solutions.problem_2 import DFT_2
import os
from solvers import Individual_3, AlphaInherit


class DFT_5(DFT_2):
    @property
    def __name__(self):
        return "Problem_5"

    def train(
        self, lr: float = 0.01, epoch: int = 1000, store_path="./result/Problem_5.pkl", is_post_train: bool = False, **kwargs
    ):
        error = 100000
        error_list = []
        A_K_list = []
        beta_list = []
        self.lr = lr
        for i in range(epoch):
            try:
                A_K_temp = [
                    self.identify(x - lr * self.grad(k))
                    for k, x in enumerate(self._A_K)
                ]
                beta_temp = self.beta - lr * self.grad_beta()
                if self.error - error > 0:
                    break
                if i % 100 == 0:
                    print(f"epoch:{i}, error:{self.error}")
                error = self.error
                self._A_K = A_K_temp
                self.beta = beta_temp
                error_list.append(self.error)
                A_K_list.append(self._A_K)
                beta_list.append(self.beta)
            except:
                raise Exception("Here! Error!")
        pickle.dump((self, error_list, A_K_list, beta_list), open(os.path.join(
                    store_path, f"K{self.K}_N{self.N}_lr{lr}_"+self.__name__+".pkl"), "wb"))
        if is_post_train:
            self.post_train(store_path=store_path, **kwargs)

    def post_train(
        self,
        store_path: str = "./result/Problem_3.pkl",
        iter_num: int = 10000,
        individual_num: int = 8,
        offstring_num: int = 8,
    ):
        """
        后期操作，在完成浮点训练之后，进行一次规划，由于规划的目标函数不可以求导，不能使用梯度下降，替代梯度的方法也不好做，这里尝试使用智能优化算法遗传算法来做。
        """
        print("*"*50)
        print("现在进入遗传算法内")
        print("*"*50)
        if not os.path.exists(os.path.join(store_path, 'problem_5')):
            os.mkdir(os.path.join(store_path, 'problem_5'))
        for q in range(1, 17):
            print(f"此时是q={q}的遗传算法")
            print("*"*50)
            M, _, _, _ = pickle.load(
                open(os.path.join(store_path, f"K{self.K}_N{self.N}_lr{self.lr}_"+self.__name__+".pkl"), "rb"))
            IS = AlphaInherit(
                individual_num,
                Individual_3,
                val=np.random.normal(2**q, 0.5, individual_num),
                sigma=np.random.normal(0.4, 0.02, individual_num),
                dft=M.dft,
                beta=M.get_beta,
                A_K=M.A_K,
                K=M.K,
                gate=-30,
                lr=1e-4
            )
            final_group = IS.train(
                epoch_num=iter_num,
                multiply_kwargs={"offstring_num": offstring_num},
                filter_kwargs={},
            )
            val = 0
            min_obj = None
            for i in final_group:
                if val > i.score(M.dft, M.A_K, M.K, M.get_beta):
                    val = i.score(M.dft, M.A_K, M.K, M.get_beta)
                    min_obj = i

            self._A_K = [self.convert_to_power_of_two(
                self.max_two(min_obj.val*A_K_i)) for A_K_i in M._A_K]
            self.alpha = min_obj.val
            print(
                f"拟合结果与真实值的差: {val}, 拟合因子: {min_obj.val}, 拟合因子的标准差: {min_obj.sigma}")
            pickle.dump(
                (self, min_obj),
                open(os.path.join(store_path, 'problem_5', f"q{q}"+self.__name__ +
                                  "_processed" + ".pkl"), "wb"),
            )
